import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GINConv, global_add_pool


class GNNEncoder(nn.Module):

    def __init__(self, in_dim, hidden_dim=128, num_layers=6, dropout=0.3):
        super().__init__()
        self.convs = nn.ModuleList()
        self.bns = nn.ModuleList()
        self.dropout = dropout

        self.convs.append(GINConv(nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )))
        self.bns.append(nn.BatchNorm1d(hidden_dim))

        for _ in range(num_layers - 2):
            self.convs.append(GINConv(nn.Sequential(
                nn.Linear(hidden_dim, hidden_dim),
                nn.BatchNorm1d(hidden_dim),
                nn.ReLU(),
                nn.Linear(hidden_dim, hidden_dim)
            )))
            self.bns.append(nn.BatchNorm1d(hidden_dim))

        self.residual = True

    def forward(self, x, edge_index):
        h = x
        for i, (conv, bn) in enumerate(zip(self.convs, self.bns)):
            h_new = conv(h, edge_index)
            h_new = bn(h_new)
            if self.residual and i > 0:  # 残差连接
                h_new += h
            h = F.relu(h_new)
            h = F.dropout(h, p=self.dropout, training=self.training)
        return h


class ProjectionHead(nn.Module):

    def __init__(self, in_dim, hidden_dim=256, out_dim=128):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, out_dim)
        )

    def forward(self, x):
        return self.mlp(x)


class GraphCL(nn.Module):

    def __init__(self, in_dim, hidden_dim=128, out_dim=2, num_layers=6, dropout=0.3,
                 projection_dim=128, proj_hidden=256):
        super().__init__()
        self.encoder = GNNEncoder(in_dim, hidden_dim, num_layers, dropout)
        self.projection = ProjectionHead(hidden_dim, proj_hidden, projection_dim)
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(p=dropout),
            nn.Linear(hidden_dim // 2, out_dim)
        )

        self.temp = 0.5

    def forward(self, x, edge_index, batch=None, mode='supervised'):
        h = self.encoder(x, edge_index)
        graph_rep = global_add_pool(h, batch) if batch is not None else h

        if mode == 'supervised':
            return self.classifier(graph_rep)
        elif mode == 'contrastive':
            return self.projection(graph_rep)
        else:
            raise ValueError("Invalid mode. Choose 'supervised' or 'contrastive'")

    def contrastive_loss(self, z1, z2):
        device = z1.device
        n = z1.size(0)

        z = torch.cat([z1, z2], dim=0)  # [2N, D]
        sim = F.cosine_similarity(z.unsqueeze(1), z.unsqueeze(0), dim=-1) / self.temp

        mask = torch.eye(2 * n, dtype=torch.bool, device=device)
        pos_mask = mask.clone()
        neg_mask = ~mask
        sim = sim.masked_fill(pos_mask, float('-inf'))
        logits = F.log_softmax(sim, dim=1)
        loss = -logits[pos_mask].sum() / (2 * n)

        return loss